import os
from typing import Any, Dict, Optional, Tuple

import numpy as np
from tqdm import tqdm

from insightface_modules.face_feature_extractor import FaceFeatureExtractor
from utils.logger_config import get_logger_config


class FaceMatching:
    def __init__(
        self,
        model_name: str = "buffalo_l",
        ctx_id: Optional[int] = None,
        det_size: Tuple[int, int] = (640, 640),
        similarity_threshold: float = 0.5,
        det_thresh: float = 0.5,
    ) -> None:
        """
        初始化人脸匹配类

        参数:
            model_name: 使用的模型名称
            ctx_id: 设备ID (0表示CPU)
            det_size: 检测尺寸
            similarity_threshold: 相似度阈值
        """
        self.model_name = model_name
        self.ctx_id = ctx_id
        self.det_size = det_size
        self._similarity_threshold = similarity_threshold
        self.similarity_threshold: float = similarity_threshold
        self.det_thresh = det_thresh

        self.logger = get_logger_config(name="face_matching").get_logger()
        self.logger.info("初始化人脸分析模型...")

        # self.app: FaceAnalysis = FaceAnalysis(name=model_name)
        # self.app.prepare(ctx_id=ctx_id, det_size=det_size)

        self.face_feature_extractor = FaceFeatureExtractor(
            model_name=model_name,
            ctx_id=ctx_id,
            det_thresh=det_thresh,
            det_size=det_size,
        )

    # 使用 property 装饰器
    @property
    def similarity_threshold(self) -> float:
        """获取相似度阈值"""
        return self._similarity_threshold

    @similarity_threshold.setter
    def similarity_threshold(self, threshold: float) -> None:
        """
        设置相似度阈值

        参数:
            threshold: 新的相似度阈值，必须在 0 到 1 之间
        """
        if 0 <= threshold <= 1:
            self._similarity_threshold = threshold
        else:
            raise ValueError("相似度阈值必须在 0 到 1 之间")

    def _calculate_similarity(
        self, embedding1: np.ndarray, embedding2: np.ndarray, metric: str = "cosine"
    ) -> float:
        """
        计算两个特征向量之间的相似度

        参数:
            embedding1: 第一个特征向量
            embedding2: 第二个特征向量
            metric: 相似度度量方法，可选值: 'cosine', 'euclidean', 'euclidean_l2'

        返回:
            相似度值
        """
        if metric == "cosine":
            # 余弦相似度: 1表示完全相同，-1表示完全相反，0表示正交
            similarity = np.dot(embedding1, embedding2) / (
                np.linalg.norm(embedding1) * np.linalg.norm(embedding2)
            )
            return float(similarity)
        elif metric == "euclidean":
            # 欧氏距离: 0表示完全相同，值越大表示越不相似
            distance = np.linalg.norm(embedding1 - embedding2)
            return float(distance)
        elif metric == "euclidean_l2":
            # L2归一化后的欧氏距离
            embedding1_l2 = embedding1 / np.linalg.norm(embedding1)
            embedding2_l2 = embedding2 / np.linalg.norm(embedding2)
            distance = np.linalg.norm(embedding1_l2 - embedding2_l2)
            return float(distance)
        else:
            # 默认使用余弦相似度
            similarity = np.dot(embedding1, embedding2) / (
                np.linalg.norm(embedding1) * np.linalg.norm(embedding2)
            )
            return float(similarity)

    def compute_similarity(
        self, embedding1: np.ndarray, embedding2: np.ndarray, metric: str = "cosine"
    ) -> float:
        """
        计算两个人脸特征向量的相似度

        参数:
            embedding1: 第一个人脸特征向量
            embedding2: 第二个人脸特征向量
            metric: 相似度度量方法，可选值: 'cosine', 'euclidean', 'euclidean_l2'

        返回:
            相似度分数 (0-1之间的浮点数，越高表示越相似)
        """
        # 计算原始相似度/距离
        similarity = self._calculate_similarity(embedding1, embedding2, metric=metric)

        # 归一化处理
        # normalized_similarity, _ = self._normalize_similarity(similarity, metric)

        return similarity

    def _normalize_similarity(
        self, similarity: float, metric: str = "cosine"
    ) -> Tuple[float, float]:
        """
        根据度量方法对相似度/距离值进行归一化处理

        参数:
            similarity: 原始相似度/距离值
            metric: 相似度度量方法

        返回:
            Tuple[float, float]: (归一化后的相似度值, 匹配率百分比)
        """
        if metric == "cosine":
            # 余弦相似度: 值越大越相似，范围[-1,1]
            # 使用更严格的归一化方法，考虑到随机向量的余弦相似度期望值约为0
            # 将[0.3, 1]映射到[0, 1]，小于0.3的值映射为0
            if similarity < 0.3:
                normalized_similarity = 0.0
            else:
                normalized_similarity = (similarity - 0.3) / 0.7
            match_rate = normalized_similarity * 100
        elif metric in ["euclidean", "euclidean_l2"]:
            # 欧氏距离: 值越小越相似
            # 根据经验值设置阈值，通常人脸特征的欧氏距离在0.5-1.5之间
            # 距离小于0.8认为是同一个人，大于1.2认为是不同的人
            if similarity > 1.2:
                normalized_similarity = 0.0
            elif similarity < 0.8:
                normalized_similarity = 1.0
            else:
                # 线性映射[0.8, 1.2]到[1, 0]
                normalized_similarity = 1.0 - (similarity - 0.8) / 0.4
            match_rate = normalized_similarity * 100
        else:
            # 默认情况，保持原值
            normalized_similarity = similarity
            match_rate = normalized_similarity * 100

        return float(normalized_similarity), float(match_rate)

    def _get_best_similarity(
        self, source_embedding, target_embeddings, metric: str = "cosine"
    ):
        best_similarity = 0
        cur_idx = 0
        for target_embedding in target_embeddings:
            similarity = self._calculate_similarity(
                source_embedding, target_embedding, metric
            )
            if similarity > best_similarity:
                best_similarity = similarity
            cur_idx += 1
            self.logger.info(
                f"第{cur_idx}组数据，相似度{similarity}，当前最大相似度{best_similarity}"
            )

        return best_similarity

    def match_two_faces(
        self,
        image1_path: str,
        image2_path: str,
        metric: str = "cosine",
        threshold: float = 0.5,
        det_thresh: float = 0.5,
    ) -> Dict[str, Any]:
        """
        匹配两张图片中的人脸

        参数:
            image1_path: 第一张图片路径
            image2_path: 第二张图片路径
            metric: 相似度度量方法，可选值: 'cosine', 'euclidean', 'euclidean_l2'
            threshold: 相似度阈值，范围0-1，越高要求越严格

        返回:
            包含匹配结果的JSON结构
        """
        # 设置置信度参数
        self.face_feature_extractor.set_prepare(det_thresh)
        # 提取第一张图片的人脸特征
        embedding1, error1 = self.face_feature_extractor._extract_face_embedding(
            image1_path
        )
        # 如果没有检测到人脸，返回错误信息
        if embedding1 is None:
            # 恢复默认置信度参数
            self.face_feature_extractor.set_prepare(self.det_thresh)
            return {
                "matched": False,
                "error": error1,
                "image1": os.path.basename(image1_path),
                "image2": os.path.basename(image2_path),
                "similarity": 0,
                "match_rate": 0,
            }

        # 提取第二张图片的人脸特征
        embedding2, error2 = self.face_feature_extractor._extract_face_embedding(
            image2_path
        )
        if embedding2 is None:
            # 恢复默认置信度参数
            self.face_feature_extractor.set_prepare(self.det_thresh)
            return {
                "matched": False,
                "error": error2,
                "image1": os.path.basename(image1_path),
                "image2": os.path.basename(image2_path),
                "similarity": 0,
                "match_rate": 0,
            }

        # 计算相似度
        similarity = self.compute_similarity(embedding1, embedding2, metric)
        match_rate = similarity * 100

        # 判断是否匹配
        is_matched = similarity >= threshold

        # 恢复默认置信度参数
        self.face_feature_extractor.set_prepare(self.det_thresh)
        return {
            "matched": is_matched,
            "image1": os.path.basename(image1_path),
            "image2": os.path.basename(image2_path),
            "similarity": similarity,
            "match_rate": match_rate,
            "metric": metric,
            "threshold": threshold,
        }

    def match_face_to_faces(
        self,
        image1_path: str,
        image2_path: str,
        metric: str = "cosine",
        threshold: float = 0.5,
        det_thresh: float = 0.5,
    ) -> Dict[str, Any]:
        """
        匹配源图片和一张多人脸图片中的人脸，只要有一个人脸匹配成功即可

        参数:
            image1_path: 第一张图片路径
            image2_path: 第二张图片路径
            metric: 相似度度量方法，可选值: 'cosine', 'euclidean', 'euclidean_l2'
            threshold: 相似度阈值，范围0-1，越高要求越严格

        返回:
            包含匹配结果的JSON结构
        """
        # 设置置信度参数
        self.face_feature_extractor.set_prepare(det_thresh)
        # 提取第一张图片的人脸特征
        embedding1, error1 = self.face_feature_extractor._extract_face_embedding(
            image1_path
        )
        if embedding1 is None:
            return {
                "matched": False,
                "error": error1,
                "image1": os.path.basename(image1_path),
                "image2": os.path.basename(image2_path),
                "similarity": 0,
                "match_rate": 0,
            }

        # 提取第二张图片的人脸特征
        embedding2, error2 = (
            self.face_feature_extractor._extract_multiple_face_embeddings(image2_path)
        )
        if embedding2 is None:
            return {
                "matched": False,
                "error": error2,
                "image1": os.path.basename(image1_path),
                "image2": os.path.basename(image2_path),
                "similarity": 0,
                "match_rate": 0,
            }

        # 计算相似度
        # 因为有可能human face detector返回多个人脸，所以需要遍历每个目标图片的每个人脸，进行匹配计算，取最高相似度
        similarity = self._get_best_similarity(embedding1, embedding2, metric)

        # 判断是否匹配，如果图中最高相似度的人脸相似度大于阈值，则认为匹配成功
        is_matched = similarity > threshold
        # 计算匹配率
        match_rate = similarity * 100

        # 返回结果
        return {
            "matched": is_matched,
            "image1": os.path.basename(image1_path),
            "image2": os.path.basename(image2_path),
            "similarity": similarity,
            "match_rate": match_rate,
            "metric": metric,
            "threshold": threshold,
        }

    def match_face_to_targets(
        self,
        source_path: str,
        target_dir: str,
        metric: str = "cosine",
        threshold: float = 0.5,
        det_thresh: float = 0.5,
    ) -> Dict[str, Any]:
        """
        匹配一张源图片与目标目录中的多张图片

        参数:
            source_path: 源图片路径
            target_dir: 目标图片目录

        返回:
            包含匹配结果的JSON结构
        """
        # 检查目录是否存在
        if not os.path.exists(source_path):
            return {"error": f"源图片不存在: {source_path}", "matches": []}

        if not os.path.exists(target_dir):
            return {"error": f"目标目录不存在: {target_dir}", "matches": []}

        # 设置置信度参数
        self.face_feature_extractor.set_prepare(det_thresh)
        # 提取源图片的人脸特征
        source_embedding, error = self.face_feature_extractor._extract_face_embedding(
            source_path
        )
        if source_embedding is None:
            return {"error": error, "matches": []}

        # 获取目标目录中的所有图片
        target_files = [
            f
            for f in os.listdir(target_dir)
            if f.lower().endswith((".png", ".jpg", ".jpeg"))
        ]

        matches = []

        # 遍历目标图片
        for target_file in tqdm(target_files):
            target_path = os.path.join(target_dir, target_file)

            # 提取目标图片中的所有人脸特征
            target_embeddings, error = (
                self.face_feature_extractor._extract_multiple_face_embeddings(
                    target_path
                )
            )

            if target_embeddings is None:
                matches.append(
                    {
                        "matched": False,
                        "error": error,
                        "target_image": target_file,
                        "similarity": 0,
                        "match_rate": 0,
                    }
                )
                continue

            self.logger.info(
                f"图片{target_path}，提取到人脸数量{len(target_embeddings)}"
            )

            # 因为有可能human face detector返回多个人脸，所以需要遍历每个目标图片的每个人脸，进行匹配计算，取最高相似度
            similarity = self._get_best_similarity(
                source_embedding, target_embeddings, metric
            )

            # 判断是否匹配，如果图中最高相似度的人脸相似度大于阈值，则认为匹配成功
            is_matched = similarity > threshold
            # 计算匹配率
            match_rate = similarity * 100

            # 添加到结果列表 - 确保所有数值都是Python原生类型
            matches.append(
                {
                    "metric": metric,
                    "threshold": threshold,
                    "matched": is_matched,
                    "target_image": target_file,
                    "best_similarity": similarity,
                    "match_rate": match_rate,
                }
            )

        # 返回结果
        return {
            "source_image": os.path.basename(source_path),
            "total_targets": len(target_files),
            "matched_count": sum(1 for m in matches if m.get("matched", False)),
            "matches": matches,
        }
